"""MCP Client Manager Module"""
import json
import time
from typing import Dict, Any

from mcp.types import CallToolResult

from .connection import ServerConnection
from .config_model import MCPConfig
from loguru import logger
import asyncio

class MCPClientManager:
    """Main client manager that coordinates all components"""

    def __init__(self, config: MCPConfig | str):
        if isinstance(config, str):
            if config.endswith('.json'):
                self.config = MCPConfig.from_json_file(config)
            else:
                config_data = json.loads(config)
                self.config = MCPConfig(**config_data)
        else:
            self.config = config

        self.connections: Dict[str, ServerConnection] = {}

        for server_type in self.config.mcpServers:
            self.connections[server_type] = None

    async def initialize(self):
        """Initialize connections to all MCP servers"""
        for server_type in self.config.mcpServers:
            self.connections[server_type] = await self.create_session(server_type)
            logger.info(f"Successfully connected to mcp server: [{server_type}]")

    async def create_session(self, server_type: str) -> ServerConnection:
        """Create a new MCP session"""
        if server_type not in self.config.mcpServers:
            logger.error(f"Invalid mcp server type: {server_type}")
            raise ValueError(f"Invalid mcp server type: {server_type}")

        server_config = self.config.mcpServers[server_type]
        connection = ServerConnection(
            name=server_type,
            config={
                "command": server_config.command,
                "args": server_config.args,
                "retries": 3,
                "retry_delay": 1.0
            }
        )

        await connection.connect()
        if not connection.is_connected():
            logger.error(f"Failed to connect to mcp server: [{server_type}]")
            raise ConnectionError(f"Failed to connect to {server_type} server")
        logger.info(f"Created new session for mcp server: [{server_type}]")
        return connection

    async def destroy_session(self, server_type: str) -> None:
        """Destroy an existing MCP session"""
        if server_type in self.connections:
            await self.connections[server_type].cleanup()
            del self.connections[server_type]
            if not self.connections:
                self.tool_manager = None
            logger.info(f"Destroyed session for mcp server: [{server_type}]")

    async def cleanup(self) -> None:
        """Cleanup all resources"""
        await asyncio.gather(*[
            connection.cleanup() for connection in self.connections.values()
        ])
        self.connections.clear()
        logger.info("Cleaned up all resources")

    def is_available(self) -> bool:
        """Check if MCPClient is available
        """
        if not self.connections:
            return False

        for connection in self.connections.values():
            if connection and connection.is_connected():
                return True

        return False

    def is_client_type(self, server_type: str) -> bool:
        if server_type not in self.config.mcpServers:
            raise ValueError(f"Invalid server type: {server_type}")

        server_config = self.config.mcpServers[server_type]
        return server_config.type == "client"

    async def get_tools(self, server_type: str) -> list:
        """Get tool list for specific server"""
        if server_type not in self.connections:
            logger.error(f"Invalid mcp server type: {server_type}")
            raise ValueError(f"Invalid server type: {server_type}")

        try:
            tools = await self.connections[server_type].tool_manager.list_tools()
            return tools
        except Exception as e:
            logger.error(f"Failed to get tools for mcp server: [{server_type}]: {str(e)}")
            # Try to reconnect and get tools again
            logger.info(f"Attempting to reconnect to mcp server: [{server_type}]")
            self.connections[server_type] = await self.create_session(server_type)
            return await self.connections[server_type].tool_manager.list_tools()

    async def get_all_tools(self) -> Dict[str, list]:
        """Get tool list for all servers"""
        all_tools = {}
        for server_type, conn in self.connections.items():
            try:
                all_tools[server_type] = await conn.tool_manager.list_tools()
            except Exception as e:
                logger.error(f"Failed to get tools for mcp server: [{server_type}]: {str(e)}")
                # Try to reconnect and get tools again
                logger.info(f"Attempting to reconnect to mcp server: [{server_type}]")
                self.connections[server_type] = await self.create_session(server_type)
                all_tools[server_type] = await self.connections[server_type].tool_manager.list_tools()
        return all_tools

    async def get_filtered_tools(self, selected_servers: Dict[str, list]) -> Dict[str, list]:
        """Get filtered tools based on selected servers and tool names

        Args:
            selected_servers: Dict where keys are server names and values are lists of tool names.
                             If tool names list is empty, all tools from that server are included.

        Returns:
            Dict[str, list]: Filtered tools by server
        """
        filtered_tools = {}

        for server_name, tool_names in selected_servers.items():
            if server_name not in self.connections:
                continue

            try:
                all_server_tools = await self.get_tools(server_name)

                if not tool_names:
                    filtered_tools[server_name] = all_server_tools
                else:
                    server_filtered_tools = []
                    for tool in all_server_tools:
                        if hasattr(tool, 'name') and tool.name in tool_names:
                            server_filtered_tools.append(tool)
                    filtered_tools[server_name] = server_filtered_tools

            except Exception as e:
                logger.error(f"Failed to get filtered tools for mcp server: [{server_name}]: {str(e)}")
                filtered_tools[server_name] = []

        return filtered_tools

    async def get_filtered_tools_str(self, selected_servers: Dict[str, list]) -> str:
        """Get filtered tools as a formatted string

        Args:
            selected_servers: Dict where keys are server names and values are lists of tool names.
                             If tool names list is empty, all tools from that server are included.

        Returns:
            str: Formatted string of filtered tools
        """
        tools_dict = await self.get_filtered_tools(selected_servers)
        return self._format_tools_dict(tools_dict)

    async def get_all_tools_str(self) -> str:
        """Get all tools as a formatted string"""
        tools_dict = await self.get_all_tools()
        return self._format_tools_dict(tools_dict)

    def _format_tools_dict(self, tools_dict: Dict[str, list]) -> str:
        """Format tools dictionary into a numbered string

        Args:
            tools_dict: Dictionary of server names to tool lists

        Returns:
            str: Formatted string with numbered tools
        """
        result = []
        tool_index = 1

        for server, tools in tools_dict.items():
            if tools:
                for tool in tools:
                    result.append(f"{tool_index}. {tool.format_for_llm()}")
                    tool_index += 1

        return "\n".join(result)

    async def execute_tool(self, server_type: str, tool_name: str, arguments: Dict[str, Any]) -> CallToolResult:
        """Execute tool on specific server"""
        if server_type not in self.connections:
            logger.error(f"Invalid server type: {server_type}")
            raise ValueError(f"Invalid server type: {server_type}")

        try:
            start_time = time.time()
            result = await self.connections[server_type].tool_manager.execute_tool(tool_name, arguments)
            execution_time = time.time() - start_time
            logger.info(f"Successfully executed mcp tool '{tool_name}' on server '{server_type}' in {execution_time:.3f}s")
            return result
        except Exception as e:
            logger.error(f"Failed to execute tool {tool_name} on mcp server: [{server_type}]: {str(e)}")
            self.connections[server_type] = await self.create_session(server_type)
            result = await self.connections[server_type].tool_manager.execute_tool(tool_name, arguments)
            return result
